#! -*- coding: utf-8 -*-
"""
@Author: AI
@Create Time: 20240625
@Info: 模型评价
"""
import torch

from arguments import args


class SpanEvaluator(object):
    """
    SpanEvaluator computes the precision, recall and F1-score for span detection.
    """

    def __init__(self):
        super().__init__()
        self.num_infer_spans = 0
        self.num_label_spans = 0
        self.num_correct_spans = 0

    def compute(self, start_probs, end_probs, gold_start_ids, gold_end_ids):
        """
        Computes the precision, recall and F1-score for span detection.
        """
        pred_start_ids = get_bool_ids_greater_than(start_probs)
        pred_end_ids = get_bool_ids_greater_than(end_probs)
        gold_start_ids = get_bool_ids_greater_than(gold_start_ids.tolist())
        gold_end_ids = get_bool_ids_greater_than(gold_end_ids.tolist())
        num_correct_spans = 0
        num_infer_spans = 0
        num_label_spans = 0

        for predict_start_ids, predict_end_ids, label_start_ids, label_end_ids in zip(
                pred_start_ids, pred_end_ids, gold_start_ids, gold_end_ids):
            [_correct, _infer, _label] = self.eval_span(predict_start_ids, predict_end_ids,
                                                            label_start_ids, label_end_ids)
            num_correct_spans += _correct
            num_infer_spans += _infer
            num_label_spans += _label

        return num_correct_spans, num_infer_spans, num_label_spans

    def update(self, num_correct_spans, num_infer_spans, num_label_spans):
        """
        This function takes (num_infer_spans, num_label_spans, num_correct_spans) as input,
        to accumulate and update the corresponding status of the SpanEvaluator object.
        """
        self.num_infer_spans += num_infer_spans
        self.num_label_spans += num_label_spans
        self.num_correct_spans += num_correct_spans

    def eval_span(self, predict_start_ids, predict_end_ids, label_start_ids, label_end_ids):
        """
        evaluate position extraction (start, end)
        return num_correct, num_infer, num_label
        input: [1, 2, 10] [4, 12] [2, 10] [4, 11]
        output: (1, 2, 2)
        """
        pred_set = get_span(predict_start_ids, predict_end_ids)     # 得到模型输出的span集合(set), e.g. {(1, 3), (4, 5)}
        label_set = get_span(label_start_ids, label_end_ids)        # 得到标签中正确的span集合(set), e.g. {(1, 3), (4, 5), (8, 9)}
        num_correct = len(pred_set & label_set)                     # 计算正确预测的span集合(两个集合求交集), e.g. {(1, 3), {4, 5}}
        num_infer = len(pred_set)
        num_label = len(label_set)
        return (num_correct, num_infer, num_label)

    def accumulate(self):
        """
        This function returns the mean precision, recall and f1 score for all accumulated minibatches.
        Returns:
            tuple: Returns tuple (`precision, recall, f1 score`).
        """
        precision = float(self.num_correct_spans / self.num_infer_spans) if self.num_infer_spans else 0.
        recall = float(self.num_correct_spans / self.num_label_spans) if self.num_label_spans else 0.
        f1_score = float(2 * precision * recall / (precision + recall)) if self.num_correct_spans else 0.
        return precision, recall, f1_score

    def reset(self):
        """
        Reset function empties the evaluation memory for previous mini-batches.
        """
        self.num_infer_spans = 0
        self.num_label_spans = 0
        self.num_correct_spans = 0

    def name(self):
        """
        Return name of metric instance.
        """
        return "precision", "recall", "f1"


def evaluate(model, data_loader):
    model.eval()
    with torch.no_grad():
        for batch in data_loader:
            start_prob, end_prob = model(input_ids=batch['input_ids'].to(args.device),
                                         token_type_ids=batch['token_type_ids'].to(args.device),
                                         attention_mask=batch['attention_mask'].to(args.device))
            start_ids = batch['start_label'].to(torch.float32).detach().numpy()
            end_ids = batch['end_label'].to(torch.float32).detach().numpy()
            num_correct, num_infer, num_label = metric.compute(start_prob.cpu().detach().numpy(),
                                                               end_prob.cpu().detach().numpy(),
                                                               start_ids,
                                                               end_ids)
            metric.update(num_correct, num_infer, num_label)

        precision, recall, f1 = metric.accumulate()